# encoding: utf-8
# 项目名称： drf-simple
# @File  : throttle.py
# @Author: SunHao
# @Desc :  限流组件
# @Date  :   13:49 2025/03/23

from rest_framework import exceptions
from rest_framework import status
from rest_framework.throttling import SimpleRateThrottle
from django.core.cache import cache as default_cache


class ThrottledException(exceptions.APIException):
    status_code = status.HTTP_429_TOO_MANY_REQUESTS
    default_code = 'throttled'


class IpThrottle(SimpleRateThrottle):
    scope = "ip"
    # 设置访问频率，例如 1分钟允许访问10次
    # 其他: s、sec、m、min、h、hour、d、day
    # THROTTLE_RATES = {"user": "10/m"}
    # 在settings.py 中也设置了,那边是全局配置
    cache = default_cache

    def get_cache_key(self, request, view):
        ident = self.get_ident(request)  # 获取请求用户ID（去request中找）
        return self.cache_format % {'scope': self.scope, 'ident': ident}
        # :1:throttle_ip_127.0.0.1

    def throttle_failure(self):
        wait = self.wait()
        detail = {
            "code": 500,
            "data": "访问频率限制",
            'detail': "需等待{}s才能访问".format(int(wait))
        }
        raise ThrottledException(detail)


class UserThrottle(SimpleRateThrottle):
    scope = "user"
    # 设置访问频率，例如 1分钟允许访问10次
    # 其他: s、sec、m、min、h、hour、d、day
    # THROTTLE_RATES = {"user": "10/m"}
    # 在settings.py 中也设置了,那边是全局配置
    cache = default_cache

    def get_cache_key(self, request, view):
        ident = request.user.pk  # 用户ID
        return self.cache_format % {'scope': self.scope, 'ident': ident}
        # :1:throttle_user_127.0.0.1

    def throttle_failure(self):
        wait = self.wait()
        detail = {
            "code": 500,
            "data": "访问频率限制",
            'detail': "需等待{}s才能访问".format(int(wait))
        }
        raise ThrottledException(detail)
